`include "aes_core_h.v"

module aes_core_datapath
(
  //OUTPUT
  output [  `COL_SIZE - 1:0] bus_rd,      //Bus read signal
  output [`COUNT_SIZE - 1:0] rd_count,    //Output of counter
  output start,                           //Start Bit
  output interrupt,                       //Interrupt Bit
  output [1:0]mode,                       //Modes Bits
  output enc_dec,                         //Encryption/Decryption operation selector
  //INPUT
  input [`COL_SIZE - 1:0] bus_wr,         //Bus write signal
  input [ `SDW_NUM - 1:0] sdwcol_en,      //Enable signal to shadow registers
  input [ `COL_NUM - 1:0] col_en,         //Enable signal to column registers
  input [ `KEY_NUM - 1:0] key_en,         //Enable signal to key registers
  input config_en,                        //Enable signal to cofiguration register
  input config_src,                       //Select signal to cofiguration register input
  input col_src,                          //Select signal to column registers input
  input key_src,                          //Select signal to key registers input
  input [1:0] bus_src,                    //Select signal to MUX_BUS
  input [2:0] sbox_src,                   //Select signal to MUX_SBOX
  input [1:0] rk_src,                     //Select signal to MUX_RK
  input [1:0] rkkey_src,                  //Select signal to MUX_KEY
  input rd_count_clear,                   //Signal to clear round counter
  input rd_count_en,                      //Enable signal to round counter
  input start_set,                        //Input to start bit
  input int_set,                          //Input to interruption bit
  input [1:0] bus_reg_src,                //Select signal to MUX_BUS
  input key_gen,                          //Indicates that core is in key generation phase
  input bus_read_en,                      //Enables Bus read operation
  input clk,                              //System Clock
  input rst_n                             //Reset, asynchrounous, active low
);

localparam INT_MASK    = `INT_MASK;
localparam INT_BIT     = `INT_BIT;
localparam START_BIT   = `START_BIT;
localparam MODE_BIT_0  = `MODE_BIT_0;
localparam MODE_BIT_1  = `MODE_BIT_1;

localparam COUNT_SIZE  = `COUNT_SIZE;
localparam CONFIG_NUM  = `CONFIG_NUM;
localparam SDW_NUM     = `SDW_NUM;
localparam COL_NUM     = `COL_NUM;
localparam KEY_NUM     = `KEY_NUM;
localparam SDW_SIZE    = `SDW_SIZE;
localparam COL_SIZE    = `COL_SIZE;
localparam KEY_SIZE    = `KEY_SIZE;
localparam CONFIG_SIZE = `CONFIG_SIZE;

localparam MUXSBOX_ROTWORD = 3'b100;

wire [SDW_SIZE - 1:0] sdwcol3_out, sdwcol2_out, sdwcol1_out, sdwcol0_out, sdwcol_in;
wire [SDW_NUM*SDW_SIZE - 1:0] sdwcol_out;

//wire sdwcol3_en, sdwcol2_en, sdwcol1_en, sdwcol0_en;

wire [COL_SIZE - 1:0] col3_out, col2_out, col1_out, col0_out;
wire [COL_SIZE - 1:0] col3_sft, col2_sft, col1_sft, col0_sft;
wire [COL_NUM*COL_SIZE - 1:0] col_out;
wire [COL_SIZE - 1:0] col_bus;

wire [KEY_NUM*KEY_SIZE - 1:0] key_out;
wire [KEY_NUM*KEY_SIZE - 1:0] key_in;
wire [KEY_SIZE - 1:0] key3_in, key2_in, key1_in, key0_in;
wire [KEY_SIZE - 1:0] key3_out, key2_out, key1_out, key0_out;

wire [COL_SIZE - 1:0] sbox_in;
wire [COL_SIZE - 1:0] mixcol_in;
wire [COL_SIZE - 1:0] mixcol_out;
wire [COL_SIZE - 1:0] rot_word;
wire [COL_SIZE - 1:0] add_rk_col;
wire [COL_SIZE - 1:0] add_rk_key;
wire [COL_SIZE - 1:0] add_rk_key_enc;
wire [COL_SIZE - 1:0] add_rk_key_dec;
wire [COL_SIZE - 1:0] bus_rd_in;
wire [COL_SIZE - 1:0] mux_config_out;

wire [CONFIG_SIZE - 1:0] config_out;
wire [CONFIG_SIZE - 1:0] config_in;
wire start_bit;
wire int_bit;
wire sbox_en_de;

assign {sdwcol3_out, sdwcol2_out, sdwcol1_out, sdwcol0_out} = sdwcol_out;
assign {key3_out, key2_out, key1_out, key0_out} = key_out;
assign key_in  = {key3_in, key2_in, key1_in, key0_in};

ff 
#(
  .DATA_WIDTH ( COL_SIZE )
)
  REG_SDW_COL[SDW_NUM - 1:0]
(
  .q     ( sdwcol_out ),
  .d     ( sdwcol_in  ),
  .en    ( sdwcol_en  ),
  .clk   ( clk        ),
  .rst_n ( rst_n      )
);

assign {col3_out, col2_out, col1_out, col0_out} = col_out;

ff_mux_2x1
#(
  .DATA_WIDTH( COL_SIZE )
)
  REG_COL[COL_NUM - 1:0]
(  
  .q       ( col_out    ),
  .in1     ( sdwcol_out ),
  .in2     ( bus_wr     ),
  .ff_en   ( col_en     ),
  .mux_sel ( col_src    ),
  .clk     ( clk        ),
  .rst_n   ( rst_n      )
);

mux_4x1
#(
  .DATA_WIDTH ( COL_SIZE )
)
  MUX_COL
(
  .out ( col_bus    ),
  .in1 ( col0_out   ),
  .in2 ( col1_out   ),
  .in3 ( col2_out   ),
  .in4 ( col3_out   ),
  .sel ( bus_src    )
);

shiftRows
#(
  .DATA_WIDTH ( COL_SIZE )
)
SHIFT_ROWS
(
  .col0_sft ( col0_sft ),
  .col1_sft ( col1_sft ),
  .col2_sft ( col2_sft ),
  .col3_sft ( col3_sft ),
  .col0     ( col0_out ),
  .col1     ( col1_out ),
  .col2     ( col2_out ),
  .col3     ( col3_out ),
  .enc_dec  ( enc_dec  )
);

mux_5x1
#(
  .DATA_WIDTH ( COL_SIZE )
)
  MUX_SBOX
(
  .out ( sbox_in  ),
  .in1 ( col0_sft ),
  .in2 ( col1_sft ),
  .in3 ( col2_sft ),
  .in4 ( col3_sft ),
  .in5 ( rot_word ),
  .sel ( sbox_src )
);

assign sbox_en_de = enc_dec | (sbox_src == MUXSBOX_ROTWORD) | key_gen;

sBox 
#(
  .SBOX_NUM ( COL_NUM   )
)
  SBOX32
(
  .sbox_out ( mixcol_in  ),
  .sbox_in  ( sbox_in    ),
  .enc_dec  ( sbox_en_de )
);
  
mixColumns MIX_COL
(
  .mix_col ( mixcol_out ),
  .col     ( mixcol_in  ),
  .enc_dec ( enc_dec    )
);
  
mux_3x1 
#(
  .DATA_WIDTH ( COL_SIZE )
)
  MUX_RK
(
  .out ( add_rk_col ),
  .in1 ( col_bus    ),
  .in2 ( mixcol_in  ),
  .in3 ( mixcol_out ),
  .sel ( rk_src     )
);
  
addRoundKey 
#(
  .DATA_WIDTH ( COL_SIZE )
)
  ADD_RK
(
  .out ( sdwcol_in  ),
  .in  ( add_rk_col ),
  .key ( add_rk_key )
);

ff_mux_2x1
#(
  .DATA_WIDTH ( KEY_SIZE )
)
  REG_KEY[KEY_NUM - 1:0]
(  
  .q       ( key_out    ),
  .in1     ( key_in     ),
  .in2     ( bus_wr     ),
  .ff_en   ( key_en     ),
  .mux_sel ( key_src    ),
  .clk     ( clk        ),
  .rst_n   ( rst_n      )
);

key_expander KEY_EXP
(
  .key0_out ( key0_in          ),
  .key1_out ( key1_in          ),
  .key2_out ( key2_in          ),
  .key3_out ( key3_in          ),
  .rot_word ( rot_word         ),
  .key0     ( key0_out         ),
  .key1     ( key1_out         ),
  .key2     ( key2_out         ),
  .key3     ( key3_out         ),
  .sbox     ( mixcol_in        ),
  .round    ( rd_count         ),
  .enc_dec  ( enc_dec | key_gen)
);

mux_4x1
#(
  .DATA_WIDTH ( KEY_SIZE )
)
  MUX_KEY
(
  .out ( add_rk_key_enc ),
  .in1 ( key0_out       ),
  .in2 ( key1_out       ),
  .in3 ( key2_out       ),
  .in4 ( key3_out       ),
  .sel ( rkkey_src      )
);

inv_mixColumns INV_MIX_COL
(
  .mix_col ( add_rk_key_dec ),
  .col     ( add_rk_key_enc )
);

assign add_rk_key = (enc_dec || (!enc_dec && (rd_count == 0 | rd_count == 10))) ? add_rk_key_enc : add_rk_key_dec;

counter_clear 
#(
  .COUNT_WIDTH( COUNT_SIZE )
)
ROUND_COUNTER
(
  .count       ( rd_count       ),
  .count_clear ( rd_count_clear ),
  .count_en    ( rd_count_en    ),
  .clk         ( clk            ),
  .rst_n       ( rst_n          )
);

//=============================================================================
// 1)  Configuration Bits
//=============================================================================
assign interrupt = int_set & config_out[INT_MASK];//config_out[INT_BIT] & config_out[INT_MASK];
assign mode      = config_out[MODE_BIT_1 : MODE_BIT_0];
assign start     = config_out[START_BIT];
assign enc_dec   = config_out[MODE_BIT_0];

assign start_bit  = config_src ? start_set : bus_wr[START_BIT];
assign int_bit    = config_src ? int_set   : (bus_wr[INT_BIT] || start_bit ? 1'b0 : config_out[INT_BIT]);

generate
  genvar i;
  for(i = 0; i < CONFIG_SIZE; i = i + 1)
    begin
      if(i == START_BIT)
        assign config_in[i] = start_bit;
      else
        if(i == INT_BIT)
          assign config_in[i] = int_bit;
        else
            assign config_in[i] = config_src ? config_out[i] : bus_wr[i];
    end
endgenerate
 
ff 
#(
  .DATA_WIDTH ( CONFIG_SIZE )
)
  REG_CONFIG[CONFIG_NUM - 1:0]
(
  .q     ( config_out ),
  .d     ( config_in  ),
  .en    ( config_en  ),
  .clk   ( clk        ),
  .rst_n ( rst_n      )
);

assign mux_config_out = {{(COL_SIZE - CONFIG_SIZE){1'b0}}, config_out};

mux_3x1 
#(
  .DATA_WIDTH ( COL_SIZE )
)
  MUX_BUS
(
  .out ( bus_rd_in       ),
  .in1 ( col_bus         ),
  .in2 ( add_rk_key      ),
  .in3 ( mux_config_out  ),
  .sel ( bus_reg_src     ) 
);

ff 
#(
  .DATA_WIDTH ( COL_SIZE )
)
  REG_BUS_READ
(
  .q     ( bus_rd      ),
  .d     ( bus_rd_in   ),
  .en    ( bus_read_en ),
  .clk   ( clk         ),
  .rst_n ( rst_n       )
);


endmodule

//-----------------------------------------------------------------------------
//                            MODULES DEFIITIONS
//-----------------------------------------------------------------------------

//=============================================================================
// 1)  Register Definition
//=============================================================================
module ff
#(
  parameter DATA_WIDTH  = 32,         //FF data width
  parameter RESET_VALUE = 0           //FF value after reset
)(
  //OUTPUTS
  output reg [DATA_WIDTH - 1:0] q,    //FF output
  //INPUTS
  input wire [DATA_WIDTH - 1:0] d,    //FF input
  input wire en,                      //enable writing in FF
  input wire clk,                     //clock
  input wire rst_n                    //Reset, asynchrounous and active low
);

always @(posedge clk, negedge rst_n)
  if(!rst_n)
    q <= RESET_VALUE;
  else if(en)
    q <= d;
      
endmodule

//=============================================================================
// 2)  Register with MUX 2x1 Definition
//     mux_sel = 1 -> in1
//             = 0 -> in2
//=============================================================================
module ff_mux_2x1
#(
  parameter DATA_WIDTH  = 32,         //FF data width
  parameter RESET_VALUE = 0           //FF value after reset
)(
  //OUTPUTS
  output reg [DATA_WIDTH - 1:0] q,    //FF output
  //INPUTS
  input wire [DATA_WIDTH - 1:0] in1,  //input to MUX
  input wire [DATA_WIDTH - 1:0] in2,  //input to MUX
  input wire ff_en,                   //enable writing in FF
  input wire mux_sel,                 //Select input to FF
  input wire clk,                     //clock
  input wire rst_n                    //reset, asynchrounous and active low
);

wire [DATA_WIDTH - 1:0] mux_out;

assign mux_out = mux_sel ? in1 : in2;

always @(posedge clk, negedge rst_n)
  if(!rst_n)
    q <= RESET_VALUE;
  else if(ff_en)
    q <= mux_out;
      
endmodule

//=============================================================================
// 3)  MUX 5x1 Definition
//=============================================================================
module mux_5x1
#(
  parameter DATA_WIDTH = 32
)(
  //OUTPUTS
  output reg  [DATA_WIDTH - 1:0] out,  //MUX Output
  //INPUTS
  input  wire [DATA_WIDTH - 1:0] in1,  //MUX input - 000
  input  wire [DATA_WIDTH - 1:0] in2,  //MUX input - 001
  input  wire [DATA_WIDTH - 1:0] in3,  //MUX input - 010
  input  wire [DATA_WIDTH - 1:0] in4,  //MUX input - 011
  input  wire [DATA_WIDTH - 1:0] in5,  //MUX input - 100
  input  wire [2:0]sel                 //MUX selector
);
  
  always @(*)
    case(sel)
      3'b000:  out = in1;
      3'b001:  out = in2;
      3'b010:  out = in3;
      3'b011:  out = in4;
      3'b100:  out = in5;
      default: out = {DATA_WIDTH{1'b0}};
    endcase
endmodule

//=============================================================================
// 4)  MUX 4x1 Definition
//=============================================================================
module mux_4x1
#(
  parameter DATA_WIDTH = 32
)(
  //OUTPUTS
  output reg  [DATA_WIDTH - 1:0] out,  //MUX Output
  //INPUTS
  input  wire [DATA_WIDTH - 1:0] in1,  //MUX input - 00
  input  wire [DATA_WIDTH - 1:0] in2,  //MUX input - 01
  input  wire [DATA_WIDTH - 1:0] in3,  //MUX input - 10
  input  wire [DATA_WIDTH - 1:0] in4,  //MUX input - 11
  input  wire [1:0]sel                 //MUX selector
);

  always @(*)
    case(sel)
      2'b00:   out = in1;
      3'b01:   out = in2;
      2'b10:   out = in3;
      2'b11:   out = in4;
      default: out = {DATA_WIDTH{1'b0}};
    endcase
endmodule

//=============================================================================
// 5)  MUX 3x1 Definition
//=============================================================================
module mux_3x1
#(
  parameter DATA_WIDTH = 32
)(
  //OUTPUTS
  output reg  [DATA_WIDTH - 1:0] out,  //MUX Output
  //INPUTS
  input  wire [DATA_WIDTH - 1:0] in1,  //MUX input - 00
  input  wire [DATA_WIDTH - 1:0] in2,  //MUX input - 01
  input  wire [DATA_WIDTH - 1:0] in3,  //MUX input - 10
  input  wire [1:0]sel                 //MUX selector
 ); 
  always @(*)
    case(sel)
      2'b00:   out = in1;
      2'b01:   out = in2;
      2'b10:   out = in3;
      default: out = {DATA_WIDTH{1'b0}};
    endcase
endmodule

//=============================================================================
// 6)  Counter with clear signal Definition
//=============================================================================
module counter_clear
#(
  parameter COUNT_WIDTH = 4 
)(
  //OUTPUTS
  output reg [COUNT_WIDTH - 1:0] count,   //Counter Output
  //INPUTS
  input count_en,                         //Enables the counter
  input count_clear,                      //Clear counter: when high the counter is restarted
  input clk,                              //Counter clock
  input rst_n                             //reset: asynchrounous, active low
);

  always @(posedge clk, negedge rst_n)
    begin
      if(!rst_n)
        count <= {COUNT_WIDTH{1'b0}};
      else
        begin
          if(count_clear)
            count <= {COUNT_WIDTH{1'b0}};
          else
            if(count_en)
              count <= count + 1;
        end
    end
endmodule
